# -*- coding: utf-8 -*-

import warnings
import socket
import time
# 在导入paramiko之前禁用所有相关警告
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", module="cryptography")
warnings.filterwarnings("ignore", module="paramiko")
warnings.filterwarnings("ignore", message=".*TripleDES.*")
warnings.filterwarnings("ignore", message=".*deprecated.*")
warnings.filterwarnings("ignore", message=".*algorithm.*")

import paramiko
import os

class SSHClient:
    def __init__(self, host, port, username, password=None, key_file=None):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.key_file = key_file
        self.client = None
        self.sftp = None
        
    def connect(self, timeout=10):
        """连接SSH服务器"""
        try:
            # 创建SSH客户端
            self.client = paramiko.SSHClient()
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 设置连接参数
            connect_kwargs = {
                'hostname': self.host,
                'port': self.port,
                'username': self.username,
                'timeout': timeout
            }
            
            # 根据认证方式设置参数
            if self.key_file and os.path.exists(self.key_file):
                # 使用密钥认证
                connect_kwargs['key_filename'] = self.key_file
            else:
                # 使用密码认证
                connect_kwargs['password'] = self.password
            
            # 尝试连接
            self.client.connect(**connect_kwargs)
            
            # 打开SFTP会话
            self.sftp = self.client.open_sftp()
            
        except socket.timeout:
            raise Exception(f"连接超时: 无法在{timeout}秒内连接到{self.host}:{self.port}")
        except paramiko.AuthenticationException:
            raise Exception("认证失败: 用户名或密码错误")
        except paramiko.SSHException as e:
            raise Exception(f"SSH协议错误: {str(e)}")
        except socket.gaierror:
            raise Exception(f"DNS解析失败: 无法解析主机名 {self.host}")
        except Exception as e:
            raise Exception(f"连接SSH服务器失败: {str(e)}")
        
    def upload_file(self, local_path, remote_path):
        """上传文件，支持断点续传"""
        if not self.sftp:
            raise Exception("SSH连接未建立")
            
        # 获取本地文件大小
        if not os.path.exists(local_path):
            raise Exception(f"本地文件不存在: {local_path}")
            
        local_size = os.path.getsize(local_path)
        
        # 检查远程是否已存在部分上传的文件
        try:
            remote_stat = self.sftp.stat(remote_path)
            remote_size = remote_stat.st_size
            
            # 如果远程文件较小，则断点续传
            if remote_size < local_size:
                with open(local_path, 'rb') as local_file:
                    local_file.seek(remote_size)
                    with self.sftp.open(remote_path, 'ab') as remote_file:
                        transferred = remote_size
                        while True:
                            data = local_file.read(32768)  # 32KB chunks
                            if not data:
                                break
                            remote_file.write(data)
                            transferred += len(data)
            else:
                # 直接上传覆盖
                self.sftp.put(local_path, remote_path)
        except IOError:
            # 远程文件不存在，直接上传
            self.sftp.put(local_path, remote_path)
            
    def download_file(self, remote_path, local_path):
        """下载文件，支持断点续传"""
        if not self.sftp:
            raise Exception("SSH连接未建立")
            
        # 检查远程文件是否存在
        try:
            remote_stat = self.sftp.stat(remote_path)
            remote_size = remote_stat.st_size
        except IOError:
            raise Exception(f"远程文件不存在: {remote_path}")
        
        # 检查本地是否已存在部分下载的文件
        if os.path.exists(local_path):
            local_size = os.path.getsize(local_path)
            
            # 如果本地文件较小，则断点续传
            if local_size < remote_size:
                with self.sftp.open(remote_path, 'rb') as remote_file:
                    remote_file.seek(local_size)
                    with open(local_path, 'ab') as local_file:
                        transferred = local_size
                        while True:
                            data = remote_file.read(32768)  # 32KB chunks
                            if not data:
                                break
                            local_file.write(data)
                            transferred += len(data)
            else:
                # 直接下载覆盖
                self.sftp.get(remote_path, local_path)
        else:
            # 本地文件不存在，直接下载
            self.sftp.get(remote_path, local_path)
            
    def close(self):
        """关闭连接"""
        try:
            if self.sftp:
                self.sftp.close()
        except:
            pass
        finally:
            self.sftp = None
            
        try:
            if self.client:
                self.client.close()
        except:
            pass
        finally:
            self.client = None
            
    def listdir(self, path):
        """列出目录内容"""
        if not self.sftp:
            raise Exception("SSH连接未建立")
        return self.sftp.listdir(path)
        
    def isdir(self, path):
        """检查路径是否为目录"""
        import stat
        if not self.sftp:
            raise Exception("SSH连接未建立")
        try:
            file_stat = self.sftp.stat(path)
            return stat.S_ISDIR(file_stat.st_mode)
        except:
            return False